#!/usr/bin/env python
# coding=utf-8

from sqlalchemy import Column, String, create_engine, Integer, or_, text, exists, not_
from sqlalchemy.orm import sessionmaker, aliased
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.sql.expression import func

Base = declarative_base()
engine = create_engine('mysql+mysqlconnector://patchdetect:patchdetect@10.141.209.138:6603/patchdetect', echo=False)
DBSession = sessionmaker(bind=engine)
session = DBSession()
class BS(Base):
    __tablename__ = "binary-source"
    id = Column(Integer, primary_key=True, autoincrement=True)
    binary = Column(String(100))
    source = Column(String(100))

    def __repr__(self):
        return "<Package(id={}, binary=\"{}\", source=\"{}\")>".format(self.id, self.binary, self.source)
class CVE(Base):
    __tablename__ = "cve"
    id = Column(String(20), primary_key=True)
    summary = Column(String(1000))
class VSL(Base):
    __tablename__ = "vulnerable-software-list"
    id = Column(Integer, primary_key=True)
    cve = Column(String(20))
    product = Column(String(100))
class CPE(Base):
    __tablename__ = "CPE"
    id = Column(Integer, primary_key=True, autoincrement=True)
    cpe = Column(String(255))
class DST(Base):
    __tablename__ = "DebianSecurityTrackerGroundTruth"
    id = Column(Integer, primary_key=True, autoincrement=True)
    package = Column(String(100))
    cve = Column(String(20))
class Result(Base):
    __tablename__ = "TEST_RESULT"
    id = Column(Integer, primary_key=True, autoincrement=True)
    cve = Column(String(20))
    text = Column(String(1000))
    result = Column(String(5000))
    version = Column(String(20))
    def __str__(self):
        return "<Result(cve={}, text=\"{}\", result=\"{}\")>".format(self.cve, self.text, self.result)

    def __repr__(self):
        return str(self)

class NodeResult(Base):
    __tablename__ = "TEST_RESULT_NODE"
    id = Column(Integer, primary_key=True, autoincrement=True)
    text = Column(String(1000))
    result = Column(String(5000))
    def __str__(self):
        return "<Result(cve={}, text=\"{}\", result=\"{}\")>".format(self.cve, self.text, self.result)

    def __repr__(self):
        return str(self)

# Base.metadata.create_all(engine)


def storeResult(cve, text, result):
    try:
        session.add(Result(text=text, result=result, version="0.15", cve=cve))
        session.commit()
    except:
        session.rollback()

def storeNodeResult(text, result):
    try:
        session.add(NodeResult(text=text, result=result))
        session.commit()
    except:
        session.rollback()

def init_cpe():
    '''
    从NVD的XML文件中导入所有的CPE数据到database中
    '''
    from lxml import etree
    with open("input/official-cpe-dictionary_v2.2.xml") as f:
        tree = etree.parse(f)
    root = tree.getroot()
    cpe_list = []
    for cpe_ele in root.getchildren()[1:]:
        cpe_list.append(dict(cpe=cpe_ele.get("name")))
    try:
        session.bulk_insert_mappings(
            CPE,
            cpe_list
        )
        session.commit()
    except:
        session.rollback()

def GetSummaryByCve(cve):
    return session.query(CVE.summary).filter(CVE.id==cve).scalar()

def GetCveList():
    return session.query(CVE.id)

def GetCpeByCve(cve):
    return session.query(VSL.product).filter(VSL.cve==cve).all()

def GetSourcePackageList():
    return session.query(BS.source).distinct(BS.source).all()

def GetCpeList():
    return session.query(CPE.cpe).all()

def InsertDstGroundTruth(pkg, cve_list):
    data = [dict(package=pkg, cve=x) for x in cve_list]
    try:
        session.bulk_insert_mappings(
            DST,
            data
        )
        session.commit()
    except:
        session.rollback()

def GetDstGT():
    return session.query(DST.package, DST.cve)

def GetErrorInTestResult():
    return session.query(Result.text, Result.result).filter(not_(Result.result.like("ROOOOOOOT%")))

def GetTestResults():
    return session.query(Result.text, Result.result).filter(Result.result.like("ROOOOOOOT%"))

def GetRandomSummary():
    return session.query(CVE.summary).order_by(func.rand())[0].summary